import openai
import os

openai.api_key=os.environ.get("OPENAI_API_KEY")

class QaRobot:
    def __init__(self,promt,num_of_round,temperature=0.0):
        self.promt=promt
        self.num_of_round=num_of_round
        self.messages=[]
        self.messages.append({"role":"system","content":promt})
        self.temperature=temperature
    
    def ask(self,question):
        try:
            self.messages.append({"role":"user","content":question})
            response=openai.ChatCompletion.create(
                model="gpt-3.5-turbo",
                messages=self.messages,
                temperature=self.temperature,
                max_tokens=2048,
                top_p=1,
            )
        except Exception as e:
            return e

        message=response["choices"][0]["message"]["content"]
        num_of_tokens=response["usage"]['total_tokens']
        
        self.messages.append({"role":"assistant","content":message})
        
        if len(self.messages)>self.num_of_round*2+1:
            del self.messages[1::]
        return message,num_of_tokens
    

import gradio as gr

cosplay="老中医"

promt= f"""你是一位${cosplay}，用中文回答病人的问题。你的回答需要满足以下条件：
1. 回答必须用中文
2. 回答限制在50个字
3. 结合本草纲目以及现代医学为用户开药方"""

qa=QaRobot(promt,10,0.0)

def answer(question,history):
    response,tks=qa.ask(question)
    return f"{response}({tks})"

demo=gr.ChatInterface(answer)

demo.launch()
